import numpy as np
import argparse
import cv2 as cv
import rospy
import serial
import math
import time
import rosgraph
import os
import sys
from sklearn.cluster import KMeans
from std_msgs.msg import Header
from sensor_msgs.msg import Image
from cv_bridge import CvBridge, CvBridgeError

rospy.init_node('rknn_control', anonymous=True)
import threading
import collections

from rknnlite.api import RKNNLite
# RKNN_MODEL = "curi.rknn"
HOST = "RK3588"

class RKNN_Agent():
    def __init__(self, args):
        self.model = args.model
        self.ttyusb = args.ttyusb
        self.bridge = CvBridge()
        self.rknn_setup()
        self.ser = serial.Serial("/dev/ttyUSB"+str(self.ttyusb), 921600)
        self.reward = args.reward

        self.subscriber = rospy.Subscriber('/usb_cam/image_raw',Image,self.callback,queue_size=1,buff_size=1000000)
        self.ec_sub = rospy.Subscriber('/encoder_data',Header,self.enc_callback)
        self.publisher = rospy.Publisher('command',Header,queue_size=1)
        self.img_pub = rospy.Publisher('rwd_img',Image,queue_size=1)

        self._feature_nums = 1000
        self._max_vel = np.pi
        self._discount_factor = 0.25
        self._infer_interval = 1 if not self.reward else 1
        self._frame_count = 0

        self._cur_action = 0.0
        self._last_action= 0.0
        self._orb = cv.ORB.create(nfeatures=1000,fastThreshold=20,scoreType=cv.ORB_FAST_SCORE)
        self._blind_mode = False
        self._max_clusters = 15
        self._input_dim = 423
        
        self._txt_file = "reward.txt"
        self._init = True
        self._row = 20
        self._col = 20
        self._grid_h = 24
        self._grid_w = 32
        self._info = True

        speed = 0
        speed_high = (int(speed ) >> 8) & 0xFF  # 取高八位
        speed_low = int(speed ) & 0xFF  # 取低八位
        data = [0xAA, 0xAF, 0x05, 0x13, 0x02, speed_high, speed_low]
        checksum = sum(data) & 0xFF
        data.append(checksum)
        self.data_bytes = bytes(data)

        self.dq = collections.deque(maxlen=4)

    def rknn_setup(self):
        self.rknn_lite = RKNNLite()

        rt = self.rknn_lite.load_rknn('/home/orangepi/tools/rknn/models/'+self.model)
        if rt!=0:
            print("Load RKNN model failed!")
            exit(rt)
        print("RKNN model loaded")

        ret = self.rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
        if ret!=0:
            print("Init runtime environment failed!")
            exit(ret)
        print("Runtime environment initialized")

    def enc_callback(self, ec_data):
        data = ec_data.seq
        self.cur_p = data
        # print(f"encoder_data: {data}")
        if  self._init and ((data>15 and data<500) or (data<4076 and data>3600)):
            self._init = False

    def callback(self,img_data):
        
        t1 = time.time()
        #print("at begin ",self._blind_mode)
        self._frame_count += 1

        # Image resize
        cv_img = self.bridge.imgmsg_to_cv2(img_data,"passthrough")
        gray_img = cv.cvtColor(cv_img, cv.COLOR_BGR2GRAY)
        # show_img = cv.resize(gray_img,(320,240))
        # cv.imshow("origin_img",show_img)
        # cv.waitKey(1)

        
        kp = self._orb.detect(gray_img)
        X = np.array([[p.pt[0],p.pt[1]] for p in kp])
        X = X.reshape(-1,2)
        
        if X.shape[0] < 10:
            centers = np.zeros((10,2))
        else:
            kmeans = KMeans(n_clusters=10,max_iter=10,n_init=1).fit(X)
            centers = kmeans.cluster_centers_
        
        obs = np.array([self._cur_action,abs(self._cur_action-self._last_action)/2.0,len(kp)/1000.0])
        for ctr in centers:
            obs = np.append(obs,ctr[0]/640.0)
            obs = np.append(obs,ctr[1]/480.0)
            
        species = np.zeros((self._row, self._col),dtype=np.float32)
        self._species = species

        for i in range(len(X)):
            p = X[i]
            species[int(p[1])//self._grid_h][int(p[0])//self._grid_w] += 1.0

        for i in range(self._row):
            for j in range(self._col):
                if len(kp)==0:
                    obs = np.append(obs,0.0)
                else:
                    obs = np.append(obs,species[i][j]/len(kp))
        
        obs = obs.astype(np.float32)
        obs = np.expand_dims(obs,0)

        # img = np.expand_dims(gray_img,0)
        # img = np.expand_dims(img,0)  # (1 , 1 , h , w ) NCHW
        
        # RKNN model interface
        if self._frame_count % self._infer_interval == 0 :
            start = time.perf_counter()
            outputs = self.rknn_lite.inference(inputs=[obs])
            end = time.perf_counter()
            dur = (end-start)*1000.0
            action = outputs[0].reshape(-1)
            self._last_action = self._cur_action
            self._cur_action = action[0]
        else:
            action = self._last_action
        
        # Serial command calculation
        speed = max(min(action[0],1),-1)
        speed = speed*(-1.0)

        # limit the range
        if self.cur_p >= 2730:
            if speed > 0:
                print(f"====={self.cur_p} exceeds range, action direction CORRECT=====")
            else:
                print(f"====={self.cur_p} exceeds range, action direction FALSE {speed}=====")
                speed = 0
        elif self.cur_p < 1366:
            if speed < 0:
                print(f"====={self.cur_p} exceeds range, action direction CORRECT=====")
            else:
                print(f"====={self.cur_p} exceeds range, action direction FALSE {speed}=====")
                speed = 0 
            
        
        speed = speed * 2 * 3.1415926536 / self._discount_factor  # 速度，单位是rad
        speed_high = (int(speed ) >> 8) & 0xFF  # 取高八位
        speed_low = int(speed ) & 0xFF  # 取低八位
        data = [0xAA, 0xAF, 0x05, 0x13, 0x02, speed_high, speed_low]
        checksum = sum(data) & 0xFF
        data.append(checksum)
        self.data_bytes = bytes(data)
        
        # Rewartd calculation
        rwd = 0.0
        if self.reward:
            #print("cal reward")
            start = time.time()
            rwd = self.rwd_cal(gray_img,kp,X,centers)
            end = time.time()
            dur = (end-start)*1000.0


        
        t2 = time.time()
        dur = (t2-t1)*1000.0
        # print("frame time: ",dur," ms")
        print("--------------------------------------")
        print("action: ",self._cur_action)
        print("reward: ",rwd)
        f =  open(self._txt_file,'a')
        f.write(str(rwd)+"\n")
        f.close()
        print("--------------------------------------")
        # self.ser.write(self.data_bytes)

    def read_data(self):
        while True:
            # data = self.ser.read(8)
            # print("I read ", data)

            data = self.ser.read(1)
            if data[0]!=2:
                self.dq.append(data[0])
                # print("data ",data)
                continue
            else:
                if(self.dq.count(170) and self.dq.count(0) and \
                   self.dq.count(218) and self.dq.count(175) and\
                    self.dq.index(170)==0 and self.dq.index(175)==1\
                    and self.dq.index(0)==2 and self.dq.index(218)==3
                    ):
                    high_byte = self.ser.read(1)[0]
                    low_byte = self.ser.read(1)[0]
                    encoder_data = high_byte*256 + low_byte
                    if int(encoder_data)>=0 and int(encoder_data)<=4096:
                        print("encoder ", int(encoder_data))

                # print("data ",data)

    def rwd_cal(self,img,kp,X,centers):
        img_reward = cv.cvtColor(img,cv.COLOR_GRAY2BGR)
        for p in kp:
            cv.circle(img_reward,(int(p.pt[0]),int(p.pt[1])),3,(0,255,0))
        for c in centers:
            cv.circle(img_reward,(int(c[0]),int(c[1])),10,(0,0,255),-1)
        if(X.size<200):
            if self._info:
                print("not enough(<200) key points detected!")
        #num features reward
        num_kp = len(kp)
        
        # weight1
        quantity_weight = min(num_kp/self._feature_nums,1.0)
        quantity_weight = quantity_weight**10

        # weight2
        species = self._species      
        species /= num_kp
        species = np.where(species<1e-10,1,species)
        species_log = np.log(species)
        score_mat = -(species_log * species)
        score_mat /= np.log(self._col*self._row)
        diversity_weight = self.function_avg(np.sum(score_mat))
        
        # if self._init:
        #     self._cur_action = 0.0
        #     self._last_action = 0.0

        # weight3
        act = abs(self._cur_action)
        act_weight = (1/(0.399*np.sqrt(2*np.pi))) * np.exp(-(act-0.5)**2/(2*0.03))
        
        # weight4
        diff = abs(self._cur_action-self._last_action)
        diff_weight = (1/(0.399*np.sqrt(2*np.pi))) * np.exp(-diff**2/(2*0.03))
        
        #显示图像

        #self.img_pub.publish(self.bridge.cv2_to_imgmsg(img_reward,"bgr8"))
        show_img = cv.resize(img_reward,(320,240))
        cv.imshow("reward image",show_img)
        cv.waitKey(1)
        
        #奖励信息
        reward = quantity_weight*diversity_weight +  act_weight*diff_weight
        return reward
    
    def close(self):
        self.ser.close()
        # self._f.close()

    def function_avg(self,x,a=0.65,b=20):
        return 0.5*((np.exp(x-a)-np.exp(-b*(x-a))) / (np.exp(x-a)+np.exp(-b*(x-a)))) + 0.5
        
if __name__=='__main__':
    parser = argparse.ArgumentParser('arg parser')
    parser.add_argument('--reward',type=bool,default=True,help='show the reward info and img')
    parser.add_argument('--ttyusb',type=int,default=0,help='the ttyusb port')
    parser.add_argument('--model',type=str,default='shannon.rknn',help='the rknn model path')
    args = parser.parse_args()
    
    agent = RKNN_Agent(args)

    rate = rospy.Rate(10)
    # tt = threading.Thread(target=agent.read_data)
    # tt.start()  

    while not rospy.is_shutdown():
        agent.ser.write(agent.data_bytes)
        rate.sleep()

    agent.close()


